import  numpy  as  np
import  matplotlib.pyplot  as  plt
from  scipy.special  import  comb

def main():
    p =  0.75    # 向右的概率
    q =  1  - p  # 向左的概率
    a =  1        # 步长
    tau =  1     # 时间间隔
    num_steps_traj = 8000       # 总步数
    num_walkers_traj_show = 10 # 显示的轨迹数量
    num_walkers_traj_stat = 2000 # 用于统计均值和方差的粒子数

    # 模拟
    steps_stat = np.random.choice([a, -a], size=(num_walkers_traj_stat, num_steps_traj), p=[p, q])
    # 计算每个粒子在每一步的位置 (累积和)
    trajectories = np.cumsum(steps_stat, axis=1)
    # 在轨迹前加上初始位置0
    trajectories = np.insert(trajectories, 0, 0, axis=1)

    # 绘制
    time_points = np.arange(num_steps_traj + 1) * tau

    # 绘制几条样本轨迹
    for i in range(num_walkers_traj_show):
        plt.plot(time_points, trajectories[i, :], alpha=0.5)

    # 计算并绘制统计均值 (来自模拟)
    mean_sim = np.mean(trajectories, axis=0)
    # 计算并绘制理论均值
    mean_theory = time_points * a * (p - q) / tau
    plt.plot(time_points, mean_theory, 'k:', linewidth=2, label=r'$\langle x \rangle$')

    # 计算并绘制理论标准差范围
    variance_theory = 4 * (time_points/tau) * (a**2) * p * q
    std_theory = np.sqrt(variance_theory)
    plt.plot(time_points, mean_theory + std_theory, 'k--', linewidth=2, label=r'$\langle x \rangle \pm \sigma$')
    plt.plot(time_points, mean_theory - std_theory, 'k--', linewidth=2)


    plt.title('Random Walk Trajectories and Moments')
    plt.xlabel('Time')
    plt.ylabel('Position')
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.6)

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()
